// This code is hosted on http://code.google.com/p/lenthorp/
// Freely available for use in applications, but should NOT be modified
// Email all comments to lenthorpresearch@gmail.com


#include <qt/processes/generalstochasticprocessarray.hpp>
#include <ql/math/matrixutilities/pseudosqrt.hpp>

namespace QuantLib {

    GeneralStochasticProcessArray::GeneralStochasticProcessArray(
        const std::vector<boost::shared_ptr<StochasticProcess> >& processes,
        const Matrix& correlation)
    : processes_(processes),
      sqrtCorrelation_(pseudoSqrt(correlation,SalvagingAlgorithm::Spectral)) {

        QL_REQUIRE(!processes.empty(), "no processes given");
		// Change to number of miniProcesses
		Size s = 0;
		for (Size i=0; i<processes.size(); ++i)
			s += processes[i]->size();
        //QL_REQUIRE(correlation.rows() == processes.size(),
		QL_REQUIRE(correlation.rows() == s,
                   "mismatch between number of processes "
                   "and size of correlation matrix");
        for (Size i=0; i<processes_.size(); i++)
            registerWith(processes_[i]);
    }

	// Change to number of miniProcesses
    Size GeneralStochasticProcessArray::size() const {
        //return processes_.size();
		Size s = 0;
		for (Size i=0; i<processes_.size(); ++i)
			s += processes_[i]->size();
		return s;
    }

    Disposable<Array> GeneralStochasticProcessArray::initialValues() const {
        Array tmp(size());
        //for (Size i=0; i<size(); ++i)
		Size count = 0;
		for (Size i=0; i<processes_.size(); ++i)
		{
			Size pSize = processes_[i]->size();
			if (pSize == 1)
				tmp[count++] = boost::static_pointer_cast<StochasticProcess1D>(processes_[i])->x0();
			else
			{
				Array tempInitValues = processes_[i]->initialValues();
				for (Size j=0; j<pSize; ++j)
					tmp[count++] = tempInitValues[j];
			}
		}
        return tmp;
    }

	Disposable<Array> GeneralStochasticProcessArray::drift(Time t,
                                                    const Array& x) const {
		return Array(0);
	}
	Disposable<Matrix> GeneralStochasticProcessArray::diffusion(
                                               Time t, const Array& x) const {
		return Matrix();
	}


    /*Disposable<Array> GeneralStochasticProcessArray::drift(Time t,
                                                    const Array& x) const {
        Array tmp(size());
        for (Size i=0; i<size(); ++i)
            tmp[i] = processes_[i]->drift(t, x[i]);
        return tmp;
    }

    Disposable<Matrix> GeneralStochasticProcessArray::diffusion(
                                               Time t, const Array& x) const {
        Matrix tmp = sqrtCorrelation_;
        for (Size i=0; i<size(); ++i) {
            Real sigma = processes_[i]->diffusion(t, x[i]);
            std::transform(tmp.row_begin(i), tmp.row_end(i),
                           tmp.row_begin(i),
                           std::bind2nd(std::multiplies<Real>(),sigma));
        }
        return tmp;
    }

    Disposable<Array> GeneralStochasticProcessArray::expectation(Time t0,
                                                          const Array& x0,
                                                          Time dt) const {
        Array tmp(size());
        for (Size i=0; i<size(); ++i)
            tmp[i] = processes_[i]->expectation(t0, x0[i], dt);
        return tmp;
    }

    Disposable<Matrix> GeneralStochasticProcessArray::stdDeviation(Time t0,
                                                            const Array& x0,
                                                            Time dt) const {
        Matrix tmp = sqrtCorrelation_;
        for (Size i=0; i<size(); ++i) {
            Real sigma = processes_[i]->stdDeviation(t0, x0[i], dt);
            std::transform(tmp.row_begin(i), tmp.row_end(i),
                           tmp.row_begin(i),
                           std::bind2nd(std::multiplies<Real>(),sigma));
        }
        return tmp;
    }*/

    Disposable<Matrix> GeneralStochasticProcessArray::covariance(Time t0,
                                                          const Array& x0,
                                                          Time dt) const {
        Matrix tmp = stdDeviation(t0, x0, dt);
        return tmp*transpose(tmp);
    }

    Disposable<Array> GeneralStochasticProcessArray::evolve(
                  Time t0, const Array& x0, Time dt, const Array& dw) const {
        const Array dz = sqrtCorrelation_ * dw;

        Array tmp(size());
		Size count = 0;
		for (Size i=0; i<processes_.size(); ++i)
		{
			Size pSize = processes_[i]->size();
			if (pSize == 1)
				tmp[count] = boost::static_pointer_cast<StochasticProcess1D>(processes_[i])->evolve(t0, x0[i], dt, dz[count]);
			else
			{
				Array tempDz(pSize);
				Array tempX0(pSize);
				for (Size j=0; j<pSize; ++j)
				{
					tempDz[j] = dz[count + j];
					tempX0[j] = x0[count + j];
				}
				Array tempValues(processes_[i]->evolve(t0, tempX0, dt, tempDz));
				for (Size j=0; j<pSize; ++j)
					tmp[count + j] = tempValues[j];
			}
			count += pSize;
		}
        return tmp;
    }

    /*Disposable<Array> GeneralStochasticProcessArray::apply(const Array& x0,
                                                    const Array& dx) const {
        Array tmp(size());
        for (Size i=0; i<size(); ++i)
            tmp[i] = processes_[i]->apply(x0[i],dx[i]);
        return tmp;
    }*/

    Time GeneralStochasticProcessArray::time(const Date& d) const {
        return processes_[0]->time(d);
    }

    const boost::shared_ptr<StochasticProcess>&
    GeneralStochasticProcessArray::process(Size i) const {
        return processes_[i];
    }

    Disposable<Matrix> GeneralStochasticProcessArray::correlation() const {
        return sqrtCorrelation_ * transpose(sqrtCorrelation_);
    }

}
